import numpy as np
from chaco.api import *
from traits.api import *
from traitsui.api import *
from chaco.tools.api import PanTool, ZoomTool
from enable.api import ComponentEditor
import time
from PyQt5.QtCore import QThread

from PyQt5.QtCore import *
from PyQt5.QtGui import *
from PyQt5.QtWidgets import *


class Form(HasTraits):
    frequencies = List(Float, [10])
    elevations = List(Float, [8])
    azimuthStart = Float(0)
    azimuthStop = Float(45)
    azimuthNumber = Int(46)
    calculatemethod=Enum('po', 'ptd', 'po+ptd')
    plot = Instance(Plot)
    progress = Str('0%')
    thread = QThread()


    def calculateall(self):
        if calculate_method == 'po':
            method_PO()
        if calculate_method == 'ptd':
            method_PTD()
        if calculate_method == 'po+ptd':
            method_POPTD()

    def method_PO(self):
        if (v := self.root.v) is None:
            return
        if (f := self.root.f) is None:
            return
        frequencies = np.array(self.frequencies)
        elevations = np.array(self.elevations)
        azimuths = np.linspace(self.azimuthStart, self.azimuthStop, self.azimuthNumber)
        self.rcs = np.full((2, frequencies.size, elevations.size, azimuths.size), np.nan)
        np.seterr(divide='ignore')
        v = np.concatenate((v[f[:, 0]][:, np.newaxis, :], v[f[:, 1]][:, np.newaxis, :], v[f[:, 2]][:, np.newaxis, :]), axis=1)
        e = np.roll(v, -1, axis=1) - v
        ec = (np.roll(v, -1, axis=1) + v) / 2
        n = np.cross(e[:, 0], e[:, 1])
        a = np.linalg.norm(n, axis=-1) / 2
        n = n / a[:, np.newaxis] / 2
        eps = 1e-8
        totalNumber = frequencies.size * elevations.size * azimuths.size
        for ff, frequency in enumerate(frequencies):
            k = 2 * np.pi / .3 * frequency
            for ee, elevation in enumerate(elevations):
                cosel = np.cos(elevation * np.pi / 180)
                sinel = np.sin(elevation * np.pi / 180)
                self.plot.title = f'Frequency={frequency}GHz,Elevation={elevation} degree'
                for aa, azimuth in enumerate(azimuths):
                    cosaz = np.cos(azimuth * np.pi / 180)
                    sinaz = np.sin(azimuth * np.pi / 180)
                    i = -np.array([cosel * cosaz, cosel * sinaz, sinel])
                    ni = n @ i
                    perpendicular = ni <= -1 + eps
                    oblique = (ni < -eps) & ~perpendicular
                    w = -2 * i
                    poPerpendicular = np.einsum('i,i,i->', ni[perpendicular], np.exp(1j * k * v[perpendicular, 0, :] @ w), a[perpendicular]) * 1j * k / np.sqrt(np.pi)
                    ncw = np.cross(n[oblique], w, axis=-1)
                    ncw2 = (ncw * ncw).sum(-1)
                    poOblique = np.einsum('i,ik,ijk,ij->', ni[oblique] / ncw2, ncw, e[oblique], np.exp(1j * k * ec[oblique] @ w) * np.sinc(k * e[oblique] @ w / 2 / np.pi)) / np.sqrt(np.pi)
                    po = poPerpendicular + poOblique
                    self.rcs[0, ff, ee, aa] = 20 * np.log10(np.abs(po))
                    self.rcs[1, ff, ee, aa] = 20 * np.log10(np.abs(po))
                    self.plot.data.update_data(
                        x=azimuths[:aa + 1],
                        hh=self.rcs[0, ff, ee, :aa + 1],
                        vv=self.rcs[1, ff, ee, :aa + 1],
                    )
                    self.progress = f'{((ff * elevations.size + ee) * azimuths.size + aa + 1) / totalNumber:<.1%}'
                    time.sleep(.1)

    def method_PTD(self):
        if (v := self.root.v) is None:
            return
        if (f := self.root.f) is None:
            return
        frequencies = np.array(self.frequencies)
        elevations = np.array(self.elevations)
        azimuths = np.linspace(self.azimuthStart, self.azimuthStop, self.azimuthNumber)
        self.rcs = np.full((2, frequencies.size, elevations.size, azimuths.size), np.nan)
        np.seterr(divide='ignore')
        eps = 1e-8
        v = np.concatenate((v[f[:, 0]][:, np.newaxis, :], v[f[:, 1]][:, np.newaxis, :], v[f[:, 2]][:, np.newaxis, :]),
                           axis=1)
        e = np.roll(v, -1, axis=1) - v
        n = np.cross(e[:, 0], e[:, 1])
        a = np.linalg.norm(n, axis=-1) / 2
        n = n / a[:, np.newaxis] / 2
        edges, edgesFacet = self.extractEdges(f)
        v = self.root.v
        et = v[edges[:, 2], :] - v[edges[:, 1], :]
        et_length = np.sqrt((et * et).sum(axis=1))
        et_unit = et / et_length[:, np.newaxis]
        ec_edge = (v[edges[:, 2], :] + v[edges[:, 1], :]) / 2
        n1 = n[edgesFacet[:, 0], :]
        n2 = n[edgesFacet[:, 1], :]
        cos_alpha = (n2 * n1).sum(axis=1)
        sin_alpha = (n2 * np.cross(et_unit, n1)).sum(axis=1)
        ne = np.arctan2(sin_alpha, cos_alpha) / np.pi + 1
        acute = ne > 1.01
        edges = edges[acute]
        edgesFacet = edgesFacet[acute]
        ec_edge = ec_edge[acute]
        et = et[acute]
        et_length = et_length[acute]
        et_unit = et_unit[acute]
        n1 = n1[acute]
        n2 = n2[acute]
        ne = ne[acute]
        totalNumber = frequencies.size * elevations.size * azimuths.size
        for ff, frequency in enumerate(frequencies):
            k = 2 * np.pi / .3 * frequency
            for ee, elevation in enumerate(elevations):
                cosel = np.cos(elevation * np.pi / 180)
                sinel = np.sin(elevation * np.pi / 180)
                self.plot.title = f'Frequency={frequency}GHz,Elevation={elevation} degree'
                for aa, azimuth in enumerate(azimuths):
                    cosaz = np.cos(azimuth * np.pi / 180)
                    sinaz = np.sin(azimuth * np.pi / 180)
                    s = np.array([cosel * cosaz, cosel * sinaz, sinel])
                    i = np.array([-cosel * cosaz, -cosel * sinaz, -sinel])
                    w = s - i
                    n1i = n1 @ i
                    n2i = n2 @ i
                    light_edge = (n1i < -eps) | (n2i < -eps)
                    n1_light = n1[light_edge]
                    n2_light = n2[light_edge]
                    ne_light = ne[light_edge]
                    et_light = et[light_edge]
                    et_unit_light = et_unit[light_edge]
                    et_length_light = et_length[light_edge]
                    ec_edge_light = ec_edge[light_edge]
                    it = et_unit_light @ i
                    it2 = it ** 2
                    xaxis = np.cross(et_unit_light, n1_light)
                    yaxis = -n1_light
                    cos_phi = xaxis @ (-i)
                    sin_phi = yaxis @ (-i)
                    phi_i1 = np.arctan2(sin_phi, cos_phi) + np.pi
                    phi_i2 = ne_light * np.pi - phi_i1
                    cos_phi = xaxis @ s
                    sin_phi = yaxis @ s
                    phi_s1 = np.arctan2(sin_phi, cos_phi) + np.pi
                    phi_s2 = ne_light * np.pi - phi_s1
                    psi_minus1 = phi_s1 - phi_i1
                    psi_minus2 = phi_s2 - phi_i2
                    psi_plus1 = phi_s1 + phi_i1
                    psi_plus2 = phi_s2 + phi_i2
                    X1 = -1 / 2 / ne_light / np.tan((np.pi - psi_minus1) / 2 / ne_light) + 1 / 2 / np.tan(
                        (np.pi - psi_minus1) / 2) * (phi_i1 < np.pi)
                    X2 = -1 / 2 / ne_light / np.tan((np.pi - psi_minus2) / 2 / ne_light) + 1 / 2 / np.tan(
                        (np.pi - psi_minus2) / 2) * (phi_i2 < np.pi)
                    Y1 = -1 / 2 / ne_light / np.tan((np.pi - psi_plus1) / 2 / ne_light) + 1 / 2 / np.tan(
                        (np.pi - psi_plus1) / 2) * (phi_i1 < np.pi)
                    Y2 = -1 / 2 / ne_light / np.tan((np.pi - psi_plus2) / 2 / ne_light) + 1 / 2 / np.tan(
                        (np.pi - psi_plus2) / 2) * (phi_i2 < np.pi)
                    X1[np.abs(np.pi - psi_minus1) < .001] = 0
                    X2[np.abs(np.pi - psi_minus2) < .001] = 0
                    Y1[np.abs(np.pi - psi_plus1) < .001] = 0
                    Y2[np.abs(np.pi - psi_plus2) < .001] = 0
                    X = X1 + X2
                    Y = Y1 + Y2
                    f = X - Y
                    g = X + Y
                    ptd_common_factor = et_length_light * np.exp(1j * k * ec_edge_light @ w) * np.sinc(
                        k * et_light @ w / 2 / np.pi) / np.sqrt(np.pi)

                    # for HH polarization
                    es = np.array([-sinaz, cosaz, 0])  ###HH
                    ei = np.array([-sinaz, cosaz, 0])
                    hs = np.cross(s, es)
                    hi = np.cross(i, es)
                    eit = et_unit_light @ ei
                    est = et_unit_light @ es
                    hit = et_unit_light @ hi
                    hst = et_unit_light @ hs
                    ptd_polarization_factor = (eit * est * f + hit * hst * g) / (1 - it2)
                    ptd = np.sum(ptd_common_factor * ptd_polarization_factor)
                    self.rcs[0, ff, ee, aa] = 20 * np.log10(np.abs(ptd))
                    # for VV polarization
                    es = np.array([-sinel * cosaz, -sinel * sinaz, cosel])  ###VV
                    ei = np.array([-sinel * cosaz, -sinel * sinaz, cosel])
                    hs = np.cross(s, es)
                    hi = np.cross(i, es)
                    eit = et_unit_light @ ei
                    est = et_unit_light @ es
                    hit = et_unit_light @ hi
                    hst = et_unit_light @ hs
                    ptd_polarization_factor = (eit * est * f + hit * hst * g) / (1 - it2)
                    ptd = np.sum(ptd_common_factor * ptd_polarization_factor)
                    self.rcs[1, ff, ee, aa] = 20 * np.log10(np.abs(ptd))
                    self.plot.data.update_data(
                        x=azimuths[:aa + 1],
                        hh=self.rcs[0, ff, ee, :aa + 1],
                        vv=self.rcs[1, ff, ee, :aa + 1],
                    )
                    self.progress = f'{((ff * elevations.size + ee) * azimuths.size + aa + 1) / totalNumber:<.1%}'
                    time.sleep(.1)

    def method_POPTD(self):
        if (v := self.root.v) is None:
            return
        if (f := self.root.f) is None:
            return
        frequencies = np.array(self.frequencies)
        elevations = np.array(self.elevations)
        azimuths = np.linspace(self.azimuthStart, self.azimuthStop, self.azimuthNumber)
        self.rcs = np.full((2, frequencies.size, elevations.size, azimuths.size), np.nan)
        np.seterr(divide='ignore')
        v = np.concatenate((v[f[:, 0]][:, np.newaxis, :], v[f[:, 1]][:, np.newaxis, :], v[f[:, 2]][:, np.newaxis, :]),
                           axis=1)
        vc = v.mean(axis=1)
        e = np.roll(v, -1, axis=1) - v
        ec = (np.roll(v, -1, axis=1) + v) / 2
        n = np.cross(e[:, 0], e[:, 1])
        a = np.linalg.norm(n, axis=-1) / 2
        n = n / a[:, np.newaxis] / 2
        edges, edgesFacet = self.extractEdges(f)
        v = self.root.v
        et = v[edges[:, 2], :] - v[edges[:, 1], :]
        et_length = np.sqrt((et * et).sum(axis=1))
        et_unit = et / et_length[:, np.newaxis]
        ec_edge = (v[edges[:, 2], :] + v[edges[:, 1], :]) / 2
        n1 = n[edgesFacet[:, 0], :]
        n2 = n[edgesFacet[:, 1], :]
        cos_alpha = (n2 * n1).sum(axis=1)
        sin_alpha = (n2 * np.cross(et_unit, n1)).sum(axis=1)
        ne = np.arctan2(sin_alpha, cos_alpha) / np.pi + 1
        acute = ne > 1.01
        edges = edges[acute]
        edgesFacet = edgesFacet[acute]
        ec_edge = ec_edge[acute]
        et = et[acute]
        et_length = et_length[acute]
        et_unit = et_unit[acute]
        n1 = n1[acute]
        n2 = n2[acute]
        ne = ne[acute]
        eps = 1e-8
        totalNumber = frequencies.size * elevations.size * azimuths.size
        for ff, frequency in enumerate(frequencies):
            k = 2 * np.pi / .3 * frequency
            for ee, elevation in enumerate(elevations):
                cosel = np.cos(elevation * np.pi / 180)
                sinel = np.sin(elevation * np.pi / 180)
                self.plot.title = f'Frequency={frequency}GHz,Elevation={elevation} degree'
                for aa, azimuth in enumerate(azimuths):
                    cosaz = np.cos(azimuth * np.pi / 180)
                    sinaz = np.sin(azimuth * np.pi / 180)
                    s = np.array([cosel * cosaz, cosel * sinaz, sinel])
                    i = np.array([-cosel * cosaz, -cosel * sinaz, -sinel])
                    w = s - i
                    ni = n @ i
                    light = ni < -eps
                    nw = np.cross(n, w, axis=-1)
                    nw2 = (nw * nw).sum(axis=-1)
                    nwSmall = nw2 < eps
                    perpendicular = nwSmall & light
                    oblique = (~nwSmall) & light
                    n1i = n1 @ i
                    n2i = n2 @ i
                    light_edge = (n1i < -eps) | (n2i < -eps)
                    n1_light = n1[light_edge]
                    n2_light = n2[light_edge]
                    ne_light = ne[light_edge]
                    et_light = et[light_edge]
                    et_unit_light = et_unit[light_edge]
                    et_length_light = et_length[light_edge]
                    ec_edge_light = ec_edge[light_edge]
                    it = et_unit_light @ i
                    it2 = it ** 2
                    xaxis = np.cross(et_unit_light, n1_light)
                    yaxis = -n1_light
                    cos_phi = xaxis @ (-i)
                    sin_phi = yaxis @ (-i)
                    phi_i1 = np.arctan2(sin_phi, cos_phi) + np.pi
                    phi_i2 = ne_light * np.pi - phi_i1
                    cos_phi = xaxis @ s
                    sin_phi = yaxis @ s
                    phi_s1 = np.arctan2(sin_phi, cos_phi) + np.pi
                    phi_s2 = ne_light * np.pi - phi_s1
                    psi_minus1 = phi_s1 - phi_i1
                    psi_minus2 = phi_s2 - phi_i2
                    psi_plus1 = phi_s1 + phi_i1
                    psi_plus2 = phi_s2 + phi_i2
                    X1 = -1 / 2 / ne_light / np.tan((np.pi - psi_minus1) / 2 / ne_light) + 1 / 2 / np.tan(
                        (np.pi - psi_minus1) / 2) * (phi_i1 < np.pi)
                    X2 = -1 / 2 / ne_light / np.tan((np.pi - psi_minus2) / 2 / ne_light) + 1 / 2 / np.tan(
                        (np.pi - psi_minus2) / 2) * (phi_i2 < np.pi)
                    Y1 = -1 / 2 / ne_light / np.tan((np.pi - psi_plus1) / 2 / ne_light) + 1 / 2 / np.tan(
                        (np.pi - psi_plus1) / 2) * (phi_i1 < np.pi)
                    Y2 = -1 / 2 / ne_light / np.tan((np.pi - psi_plus2) / 2 / ne_light) + 1 / 2 / np.tan(
                        (np.pi - psi_plus2) / 2) * (phi_i2 < np.pi)
                    X1[np.abs(np.pi - psi_minus1) < .001] = 0
                    X2[np.abs(np.pi - psi_minus2) < .001] = 0
                    Y1[np.abs(np.pi - psi_plus1) < .001] = 0
                    Y2[np.abs(np.pi - psi_plus2) < .001] = 0
                    X = X1 + X2
                    Y = Y1 + Y2
                    f = X - Y
                    g = X + Y
                    ptd_common_factor = et_length_light * np.exp(1j * k * ec_edge_light @ w) * np.sinc(
                        k * et_light @ w / 2 / np.pi) / np.sqrt(np.pi)

                    # for HH polarization
                    es = np.array([-sinaz, cosaz, 0])  ###HH
                    ei = np.array([-sinaz, cosaz, 0])
                    hs = np.cross(s, es)
                    hi = np.cross(i, es)
                    n_es_hi = n @ np.cross(es, hi)
                    poPerpendicular = np.sum(n_es_hi[perpendicular] * np.exp(1j * k * vc[perpendicular, :] @ w) * a[
                        perpendicular]) * 1j * k / np.sqrt(np.pi)
                    poOblique = np.einsum('i,ik,ijk,ij->', n_es_hi[oblique] / nw2[oblique], nw[oblique], e[oblique],
                                          np.exp(1j * k * ec[oblique] @ w) * np.sinc(
                                              k * e[oblique] @ w / 2 / np.pi)) / np.sqrt(np.pi)
                    po = poPerpendicular + poOblique
                    eit = et_unit_light @ ei
                    est = et_unit_light @ es
                    hit = et_unit_light @ hi
                    hst = et_unit_light @ hs
                    ptd_polarization_factor = (eit * est * f + hit * hst * g) / (1 - it2)
                    ptd = np.sum(ptd_common_factor * ptd_polarization_factor)

                    self.rcs[0, ff, ee, aa] = 20 * np.log10(np.abs(po + ptd))
                    # for VV polarization
                    es = np.array([-sinel * cosaz, -sinel * sinaz, cosel])  ###VV
                    ei = np.array([-sinel * cosaz, -sinel * sinaz, cosel])
                    hs = np.cross(s, es)
                    hi = np.cross(i, es)
                    n_es_hi = n @ np.cross(es, hi)
                    poPerpendicular = np.sum(n_es_hi[perpendicular] * np.exp(1j * k * vc[perpendicular, :] @ w) * a[
                        perpendicular]) * 1j * k / np.sqrt(np.pi)
                    poOblique = np.einsum('i,ik,ijk,ij->', n_es_hi[oblique] / nw2[oblique], nw[oblique], e[oblique],
                                          np.exp(1j * k * ec[oblique] @ w) * np.sinc(
                                              k * e[oblique] @ w / 2 / np.pi)) / np.sqrt(np.pi)
                    po = poPerpendicular + poOblique
                    eit = et_unit_light @ ei
                    est = et_unit_light @ es
                    hit = et_unit_light @ hi
                    hst = et_unit_light @ hs
                    ptd_polarization_factor = (eit * est * f + hit * hst * g) / (1 - it2)
                    ptd = np.sum(ptd_common_factor * ptd_polarization_factor)
                    self.rcs[1, ff, ee, aa] = 20 * np.log10(np.abs(po + ptd))
                    self.plot.data.update_data(
                        x=azimuths[:aa + 1],
                        hh=self.rcs[0, ff, ee, :aa + 1],
                        vv=self.rcs[1, ff, ee, :aa + 1],
                    )
                    self.progress = f'{((ff * elevations.size + ee) * azimuths.size + aa + 1) / totalNumber:<.1%}'
                    time.sleep(.1)

    def traits_init(self):
        self.thread.run = self.calculateall
        self.plot = Plot(ArrayPlotData())
        self.plot.data.update_data(
            x=[],
            hh=[],
            vv=[],
        )
        self.plot.title = 'Ground'
        self.plot.x_axis.title = "Azimuth/deg"
        self.plot.y_axis.title = "RCS/dBsm"
        self.plot.plot(("x", "hh"), name='HH', color='auto')
        self.plot.plot(("x", "vv"), name='VV', color='auto')
        legend = Legend()
        legend.plots = self.plot.plots
        self.plot.overlays.append(legend)

    parameter_view = View(
        Group(
            UItem('frequencies'),
            label='频率/GHz',
            show_border=True,
        ),
        Group(
            UItem('elevations'),
            label='高低角/°',
            show_border=True,
        ),
        Item('azimuthStart', label='方位角起始/°'),
        Item('azimuthStop', label='方位角终止/°'),
        Item('azimuthNumber', label='方位角数量'),
        height=600,
        resizable=True,
        buttons=OKCancelButtons,
    )
    traits_view = View(
        UItem(
            "plot",
            editor=ComponentEditor(),
            resizable=True,
        ),
        HGroup(
            UItem('progress'),
            spring
        ),
        width=1200,
        height=900,
        resizable=True,
    )

    azimuths_view = View(
        Item('azimuthStart', label='方位角起始/°'),
        Item('azimuthStop', label='方位角终止/°'),
        Item('azimuthNumber', label='方位角数量'),
        resizable=True,
        buttons=OKCancelButtons,
    )

    frequencies_view = View(
        Group(
            UItem('frequencies'),
            label='频率/GHz',
            show_border=True,
        ),
        height=600,
        resizable=True,
        buttons=OKCancelButtons,
    )

    elevations_view = View(
        Group(
            UItem('elevations'),
            label='高低角（需要大于0°/°',
            show_border=True,
        ),
        height=600,
        resizable=True,
        buttons=OKCancelButtons,
    )


    calculate_view = View(
        Item('calculatemethod', label='计算方法'),
        buttons=OKCancelButtons,
    )

